import os
from dataclasses import dataclass

from dymola_res_parser import parse_dymola_results, DymolaResult
from zq_res_parser import parse_zq_results, ZQResult


def find_common_models(dymola_dir, zq_dir):
    dymola_models = {f.rsplit('_', 1)[0] for f in os.listdir(dymola_dir) if f.endswith('.txt')}
    zq_models = {f.replace('.txt', '') for f in os.listdir(zq_dir) if f.endswith('.txt')}
    return dymola_models.intersection(zq_models)


def get_dymola_res(dymola_dir, common_models):
    retval = []
    for model in common_models:
        dymola_file = os.path.join(dymola_dir, f"{model}_500.txt")
        if os.path.exists(dymola_file):
            retval.append(parse_dymola_results(dymola_file))
    return [r for r in retval if r is not None and r.model_name is not None and r.f_evaluations is not None]


def get_zq_res(zq_dir, common_models):
    retval = []
    for model in common_models:
        zq_file = os.path.join(zq_dir, f"{model}.txt")
        if os.path.exists(zq_file):
            retval.append(parse_zq_results(zq_file))
    return [r for r in retval if r is not None and r.model_name is not None and r.f_evaluations is not None]


# Example usage
if __name__ == "__main__":
    dymola_directory = r"D:\pyProject\dymola_test\dymola_test\500-500\dymola_res"
    zq_directory = r"D:\pyProject\dymola_test\dymola_test\500-500\zq_res"
    zq_n_directory = r"D:\pyProject\dymola_test\dymola_test\500-500\zq_n_res"

    common_models = find_common_models(dymola_directory, zq_directory)
    dymola_results = get_dymola_res(dymola_directory, common_models)
    zq_results = get_zq_res(zq_directory, common_models)
    zq_n_results = get_zq_res(zq_n_directory, common_models)
    print(f"Common models: {len(common_models)}")
